import torch

from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import MNIST
import matplotlib.pyplot as plt


class MyNet(torch.nn.Module):
    """
    定义神经网络模型 MyNet:
        继承自 torch.nn.Module，实现了一个全连接神经网络。
        网络包含四个全连接层 (Linear)，其中最后一层的输出对应 MNIST 数据集的 10 个数字类别。
        forward 方法定义了前向传播的过程，使用 ReLU 作为激活函数，并在最后一层使用 log_softmax。
    """

    def __init__(self):
        super().__init__()
        # 第一层是 28 * 28 个数据
        self.fc1 = torch.nn.Linear(28 * 28, 64)  # 输入层，输入为28*28=784，输出为64
        self.fc2 = torch.nn.Linear(64, 64)  # 第二层，输入和输出均为64
        self.fc3 = torch.nn.Linear(64, 64)  # 第三层，输入和输出均为64
        self.fc4 = torch.nn.Linear(64, 10)  # 输出层，输入为64，输出为10（对应0-9的数字）

    def forward(self, x):
        x = torch.nn.functional.relu(self.fc1(x))  # 经过第一层并应用 ReLU 激活函数
        x = torch.nn.functional.relu(self.fc2(x))  # 经过第二层并应用 ReLU 激活函数
        x = torch.nn.functional.relu(self.fc3(x))  # 经过第三层并应用 ReLU 激活函数
        x = torch.nn.functional.log_softmax(self.fc4(x), dim=1)  # 经过输出层并应用 log_softmax 激活函数
        return x  # 返回输出


def get_data_loader(is_train):
    """
    定义 get_data_loader 函数:
        该函数根据 is_train 参数返回训练或测试数据加载器。
        使用 transforms.ToTensor() 将图像数据转换为 PyTorch tensor。
        使用 DataLoader 加载数据集，设置批大小为 15，并进行随机打乱。
    :param is_train:
    :return:
    """

    to_tensor = transforms.Compose([transforms.ToTensor()])  # 定义数据转换，将图片转换为 tensor

    data_set = MNIST("", is_train, transform=to_tensor, download=True)  # 下载 MNIST 数据集

    return DataLoader(data_set, batch_size=15, shuffle=True)  # 返回数据加载器，设置 batch_size 为 15，并进行随机打乱


def evaluate(test_data, net: MyNet):
    """
    定义 evaluate 函数:
        该函数用于评估模型在测试集上的准确率。
        使用 torch.no_grad() 禁用梯度计算以节省内存和计算。
        遍历测试数据集，计算预测结果并与真实标签比较，更新正确预测和总预测数量。
        返回准确率。
    :param test_data:
    :param net:
    :return:
    """

    n_correct = 0  # 记录正确预测的数量
    n_total = 0  # 记录总预测的数量

    with torch.no_grad():  # 在评估时，不计算梯度
        for (x, y) in test_data:
            outputs = net.forward(x.view(-1, 28 * 28))  # 将输入图像展开为 784 维向量
            for i, output in enumerate(outputs):
                if torch.argmax(output) == y[i]:  # 判断预测是否正确
                    n_correct += 1  # 正确预测数量加一
                n_total += 1  # 总预测数量加一
    return n_correct / n_total  # 返回准确率


def main():
    """
    定义 main 函数:
        首先获取训练和测试数据加载器。
        实例化神经网络 net。
        打印初始准确率。
        使用 Adam 优化器进行训练，训练 2 个周期，每个周期内遍历所有训练数据。
        在每个周期结束时评估模型并打印准确率。
        在最后，遍历测试数据的前四个样本，进行预测并展示每个样本的图像及其预测结果。
    :return:
    """

    train_data = get_data_loader(is_train=True)  # 获取训练数据加载器
    test_data = get_data_loader(is_train=False)  # 获取测试数据加载器

    net = MyNet()  # 实例化神经网络

    print("initial accuracy:", evaluate(test_data, net))  # 评估并打印初始准确率

    optimizer = torch.optim.Adam(net.parameters(), lr=0.001)  # 定义 Adam 优化器，学习率为 0.001

    for epoch in range(2):  # 训练两个周期
        for (x, y) in train_data:
            net.zero_grad()  # 清除之前的梯度
            output = net.forward(x.view(-1, 28 * 28))  # 进行前向传播
            loss = torch.nn.functional.nll_loss(output, y)  # 计算负对数似然损失
            loss.backward()  # 反向传播计算梯度
            optimizer.step()  # 更新参数

        print("epoch", epoch, "accuracy:", evaluate(test_data, net))  # 每个周期结束后评估并打印准确率

    for (n, (x, _)) in enumerate(test_data):  # 遍历测试数据并展示预测结果
        if n > 3:  # 只显示前四个预测结果
            break

        predict = torch.argmax(net.forward(x[0].view(-1, 28 * 28)))  # 预测第一个样本
        plt.figure(n)  # 创建新的图形
        plt.imshow(x[0].view(28, 28), cmap='gray')  # 显示图像
        plt.title("prediction: " + str(int(predict)))  # 设置标题为预测结果

    plt.show()  # 显示所有图形


if __name__ == "__main__":
    """
    整个代码实现了一个简单的全连接神经网络，用于对 MNIST 手写数字数据集进行分类。主要包括：
        数据加载与预处理。
        网络定义与前向传播。
        模型训练与评估。
        预测结果的可视化展示。
    """

    main()
